# cython: language_level=3

import random
import time
import numpy as np
cimport numpy as np
import math
from libc.stdlib cimport malloc, free

'''
简易的搜索匹配点云方法
'''

# ==========真实位置=========== #

def get_ob(int max_sample=50, double x_len=3.0, double width=2.0, double length=1.0, bint noisy=False, bint fake_point=False, double noisy_seed=0.05):
    horizon_wall = []
    vertical_wall = []
    
    def rd(seed):
        return random.uniform(-1,1)*seed

    if not noisy: noisy_seed = 0
    for i in range(max_sample):
        horizon_wall.append([-x_len-width+i*x_len/max_sample+rd(noisy_seed), 0+rd(noisy_seed), 1.0])
        vertical_wall.append([-width+rd(noisy_seed), i*length/max_sample+rd(noisy_seed), 1.0])

    _world = np.vstack(
        [np.array(horizon_wall),
        np.array(vertical_wall)]
        )

    mirror = np.vstack([
        -1*_world[:, 0],
        _world[:, 1],
        _world[:, 2]
    ]).T

    _world = np.vstack(
        [_world, mirror]
    )

    _fake = []
    for i in range(20):
        _fake.append([rd(1.0), 3.0, 1.0])

    if fake_point:
        _world = np.vstack([
            _world, np.array(_fake)
        ])

    return _world.T


# ========采样值========== #

def get_trans_matrix(double dx, double dy):
    return np.array([
        [1.0, 0.0, dx],
        [0.0, 1.0, dy],
        [0.0, 0.0, 1.0]
    ])

def get_rot_matrix(double dth):
    return np.array([
        [math.cos(dth), -math.sin(dth), 0.0],
        [math.sin(dth), math.cos(dth), 0.0],
        [0.0, 0.0, 1.0]
    ])

cdef double single_score(np.ndarray[double, ndim=1] s, np.ndarray[double, ndim=2] real):
    '''
    返回单个点的匹配度
    '''
    cdef double min_dis = math.sqrt((s[0]-real[0, 0])**2 + (s[1]-real[1, 0])**2)
    cdef int i
    cdef Py_ssize_t length = len(real[0])
    for i in range(length):
        min_dis = min(min_dis, math.sqrt((real[0,i]-s[0])**2+(real[1,i]-s[1])**2))
    return min_dis

cdef double c_match_score(np.ndarray[double, ndim=2] scan, np.ndarray[double, ndim=2] real):
    '''
    得到真实real世界和扫描scan世界的匹配系数
    '''
    cdef double score = 0.0
    cdef int i
    cdef Py_ssize_t length = len(scan[0])
    for i in range(length):
        score += single_score(np.array([scan[0,i], scan[1,i]]), real)
    # 返回每个点匹配度的平均值
    return score/length

# 类型定义
ctypedef np.ndarray ExploreStep_T(np.ndarray[double, ndim=2], np.ndarray[double, ndim=2], double*, double)

cdef np.ndarray one_x_step(
        np.ndarray[double, ndim=2] origin_scan, 
        np.ndarray[double, ndim=2] real, 
        double* new_score, 
        double step
    ):
    cdef np.ndarray[double, ndim=2] new_scan
    cdef np.ndarray[double, ndim=2] array_step = np.array([[step], [0.0], [0.0]])
    new_scan = origin_scan + array_step
    new_score[0] = c_match_score(new_scan, real)
    return new_scan

cdef np.ndarray one_y_step(
        np.ndarray[double, ndim=2] origin_scan, 
        np.ndarray[double, ndim=2] real, 
        double* new_score, 
        double step
    ):
    cdef np.ndarray[double, ndim=2] new_scan
    cdef np.ndarray[double, ndim=2] array_step = np.array([[0.0], [step], [0.0]])
    new_scan = origin_scan + array_step
    new_score[0] = c_match_score(new_scan, real)
    return new_scan

cdef np.ndarray one_th_step(
        np.ndarray[double, ndim=2] origin_scan, 
        np.ndarray[double, ndim=2] real, 
        double* new_score, 
        double step
    ):
    cdef np.ndarray[double, ndim=2] new_scan
    cdef np.ndarray[double, ndim=2] trans_matrix = np.array([
            [math.cos(step), -math.sin(step), 0.0],
            [math.sin(step), math.cos(step), 0.0],
            [0.0, 0.0, 1.0]
        ])
    new_scan = trans_matrix.dot(origin_scan)
    new_score[0] = c_match_score(new_scan, real)
    return new_scan

cdef list c_icp(np.ndarray[double, ndim=2] scan, np.ndarray[double, ndim=2] real):
    cdef:
        double new_score = 0.0 
        double origin_score = 0.0
        double x_step = 0.0
        double best_x = 0.0
        double best_y = 0.0
        double best_th = 0.0

    cdef np.ndarray origin_scan = scan
    cdef np.ndarray[double, ndim=2] new_scan = np.array([[]])

    origin_score = c_match_score(origin_scan, real)

    origin_scan = explore(one_x_step, origin_scan, real, origin_score, &new_score, 0.01, &best_x)
    origin_score = new_score

    origin_scan = explore(one_y_step, origin_scan, real, origin_score, &new_score, 0.01, &best_y)
    origin_score = new_score

    origin_scan = explore(one_th_step, origin_scan, real, origin_score, &new_score, math.radians(0.1), &best_th)
    origin_score = new_score

    cdef list ret = [best_x, best_y, best_th]
    return ret

cdef np.ndarray explore(
        ExploreStep_T func, 
        np.ndarray[double, ndim=2] origin_scan, 
        np.ndarray[double, ndim=2] real, 
        double origin_score, 
        double* new_score_ptr, 
        double step,
        double* best_step
    ):
    cdef:
        np.ndarray new_scan
        double best = 0.0
        double auto_step = 0.0
    
    # 往前试探一步
    new_scan = func(origin_scan, real, new_score_ptr, step)
    # 发现往前走增大评价值
    if new_score_ptr[0] > origin_score: 
        step = -step
        # 说明退后更好

    auto_step = step
    while True:
        # 按照理论上更好的方向变换
        new_scan = func(origin_scan, real, new_score_ptr, auto_step)
        if new_score_ptr[0] > origin_score:
            # 如果还是更差， 说明一开始在极值点，还是一开始的评分更好
            new_score_ptr[0] = origin_score
            break
        else:
            auto_step = auto_step * 1.1
            # 新的变换更优了，替换掉老的变换
            origin_scan = new_scan
            origin_score = new_score_ptr[0]
            best += auto_step

    best_step[0] = best
    return origin_scan

def final_score(np.ndarray[double, ndim=2] scan, np.ndarray[double, ndim=2] real, double tolerance = 0.01):
    '''
    获取满足tolerance要求的匹配点的数量
    '''
    cdef int amount = 0
    cdef int i
    cdef double score = 0.0
    cdef Py_ssize_t length = len(scan[0])
    for i in range(length):
        score = single_score(np.array([scan[0,i], scan[1,i]]), real)
        if abs(score) < tolerance:
            amount += 1

    return amount


''' For Python Exports '''

def icp(np.ndarray[double, ndim=2] scan, np.ndarray[double, ndim=2] real):
    return c_icp(scan, real)

def match_score(np.ndarray[double, ndim=2] scan, np.ndarray[double, ndim=2] real):
    return c_match_score(scan, real)

